#!/usr/bin/env python
"""
:mod:`pynchon.bio.alg.atom`
===========================

Algorithms on coordinates.
"""
import numpy as np
from pynchon.math.array import normalize_vector

def center_of_mass(atom, weights =None):
    """
    Calculates the (optionally weighted) center of mass of an ``AtomVector`` or
    rank-2 array of coordinates. Returns a rank-1 array.
    
    Arguments:
    
        - atom (``AtomVector`` or ``numpy.ndarray``). This can be either an
          ``AtomVector`` or 64-bit floating point rank-2 array of coordinates.  
        - weights (``sequence`` or ``numpy.ndarray``) [default: ``None``] The 
          length of the weights vector should equal the number of atoms in the 
          atom vector.
    """
    # standard atom input
    atoms = atom.asarray() if hasattr(atom, 'asarray') else atom
    if len(atoms.shape) < 2:
        raise ValueError('Expected a rank-2 array or ``AtomVector``.') 
    return np.average(atoms, axis=0, weights=weights)
    
def moments_of_inertia(atom, weights =None):
    """
    Calculates moments of inertia tensor for a vector of Atoms or a rank-2 array
    of  coordinates. Returns the tensor as a numpy 3x3 array.
    
    Arguments:
    
        - atom (``AtomVector or ``numpy.ndarray``) This can be either an
          ``AtomVector`` or 64-bit floating point rank-2 array of coordinates.
        - weights (``sequence`` or ``numpy.ndarray``) [default: ``None``] The 
          length of the weights vector should equal the number of atoms in the 
          atom vector.
    """
    # standard atom input
    atoms = atom.asarray() if hasattr(atom, 'asarray') else atom
    com = center_of_mass(atoms, weights)
    # translate to the origin
    atoms -= com
    # weights is a rank-1 array
    if weights is None:
        weights = np.ones(len(atoms))
        
    weights = np.atleast_2d(weights).transpose()
    # column code
    w_atom_diag = weights * atoms**2
    x = atoms[:,0:1]
    y = atoms[:,1:2]
    z = atoms[:,2:3]
    w_atom_xy = -weights * (x*y)
    w_atom_xz = -weights * (x*z)
    w_atom_yz = -weights * (y*z)
    I11 = np.sum(w_atom_diag[:,1:2] + w_atom_diag[:,2:3])
    I22 = np.sum(w_atom_diag[:,0:1] + w_atom_diag[:,2:3])
    I33 = np.sum(w_atom_diag[:,0:1] + w_atom_diag[:,1:2])
    I12 = np.sum(w_atom_xy)
    I13 = np.sum(w_atom_xz)
    I23 = np.sum(w_atom_yz)
    I = np.array([[I11, I12, I13],
                  [I12, I22, I23],
                  [I13, I23, I33]])
    return I

def principal_axes(atom, weights=None, sorted=True, right_handed=True):
    """
    Calculates the principal axes of a ``AtomVector``. Returns the eigenvalues
    and matrix of eigenvectors, which is a rotation matrix.
    
    Arguments:
    
        - atom (``AtomVector`` or ``numpy.ndarray``) 
        - weights (``sequence`` or ``numpy.ndarray``) 
        - sorted (``bool``) [default: ``True``] If ``True`` the eigenvectors are 
          sorted by the eigenvalues, such that the z-axis corresponds to the 
          highest eigenvalue.
        - right_handed (``bool``) [default: ``True``]
    """
    I = moments_of_inertia(atom, weights)
    vals, vects = np.linalg.eig(I)
    
    if sorted:
        idx = vals.argsort()[::-1]
        vals = vals[idx]
        vects = vects[:,idx]
    
    crs = normalize_vector(np.cross(vects[:,0],vects[:,1]))
    if right_handed and not np.allclose(crs, vects[:,2]):
        vects[:,2] = -vects[:,2]
        
    return vals, vects

def normalize_position(atom, weights=None):
    """
    
    """
    ev, rmx = principal_axes(atom, weights)
    print np.dot(atom.asarray(), rmx)
    return np.dot(atom.asarray(), rmx)
    